module mod_upward
  integer::IMAX=30,nexpand=8
  real(8)::equivposition=1.01     ! need further check h+deltah < z0
  real(8)::fftratio
contains

  subroutine init_upward(dx,dy,zsurface)
    use mod_para
    use mod_grid
    use mod_bivar
    use mod_workspace
    use mod_fftw
    use mod_interp2
    implicit none
    integer::n,i,j,k
    real(8),intent(in)::zsurface(nx2-nx1+3,ny2-ny1+3) 
    real(8)::zmin,zmax,temp,z0,dx,dy,xmin,ymin,xmax,ymax

    xmin=minval(x(:,:,nz1))
    xmax=maxval(x(:,:,nz1))
    ymin=minval(y(:,:,nz1))
    ymax=maxval(y(:,:,nz1))

    ndp=(nx2-nx1+3)*(ny2-ny1+3)
    nxi=ceiling(xmax-xmin)/dx
    nyi=ceiling(ymax-ymin)/dy
    dx=(xmax-xmin)/nxi
    dy=(ymax-ymin)/nyi

    call init_bivar()
    call init_upward_workspace(nexpand)
    do i=1,nxi
       xi(i)=xmin+dble(i-1)*dx
    enddo
    do i=1,nyi
       yi(i)=ymin+dble(i-1)*dy
    enddo

    call griddata( x(nx1-1,ny1-1,nz1), y(nx1-1,ny1-1,nz1), z(nx1-1,ny1-1,nz1), xi, yi, zi)    

    !call bivar_init2()
    call init_interp2(nxi,nyi,xi,yi,ndp,x(nx1-1,ny1-1,nz1),y(nx1-1,ny1-1,nz1))

    call dfftw_init(nxi,nyi)
    fftratio=1.D0/dble(nxi*nyi)
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    do i=1,nxi/2
       kx(i)=dble(i-1)
    enddo
    !kx(nxi/2+1)=0
    do i=nxi/2+1,nxi
       kx(i)=dble(i-nxi-1)
    enddo
    do i=1,nyi/2
       ky(i)=dble(i-1)
    enddo
    !ky(nyi/2+1)=0
    do i=nyi/2+1,nyi
       ky(i)=dble(i-nyi-1)
    enddo
    kx=2D0*pi/(dble(nxi)*dx)*kx
    ky=2D0*pi/(dble(nyi)*dy)*ky

    do j=1,nyi
       do i=1,nxi
          kk(i,j)=sqrt(kx(i)*kx(i)+ky(j)*ky(j))
          if(kk(i,j) .eq. 0) then
             kk(i,j)=epsilon(1D0)
          endif
       enddo
    enddo

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    zmin=minval(zsurface)
    zmax=maxval(zsurface)+deltah


    hnr_nfac(:,:,1)=((zmin+zmax)/2D0-zi-deltah)
    hnr_nfacdh(:,:,1)=hnr_nfac(:,:,1) +deltah
    z0=(zmax-zmin)/2D0*equivposition

    expkz0kn(:,:,0)=exp(-kk(1:nxi/2+1,:)*z0)
    do n=2,nexpand
       do j=1,nyi
          do i=1,nxi
             hnr_nfac(i,j,n)=hnr_nfac(i,j,n-1)*hnr_nfac(i,j,1)/dble(n)
             hnr_nfacdh(i,j,n)=hnr_nfacdh(i,j,n-1)*hnr_nfacdh(i,j,1)/dble(n)
          enddo
       enddo
    enddo

	hnr_nfac=hnr_nfac*fftratio
	hnr_nfacdh=hnr_nfacdh*fftratio
    do n=1,nexpand
       do j=1,nyi
          do i=1,nxi/2+1
             expkz0kn(i,j,n)=expkz0kn(i,j,n-1)*(-kk(i,j))
          enddo
       enddo
    enddo
  endsubroutine init_upward

  subroutine equivsource(u,s)

    ! workspace and fft already initilized
    ! always call fft with fftreal and fftcomplex

    use mod_para
    use mod_grid
    use mod_bivar
    use mod_workspace
    use mod_fftw
    implicit none

    
    real(8)::u(nxi,nyi) 
    real(8)::s(nxi,nyi) !equivalant source

    integer::i,j,n,iter
    integer,save::run=0



    !begin iterate

    if (run .lt. 10) then
	   s=u    
       fftreal=u
       run=run+1
	   IMAX=100
    else
       fftreal=s
	   IMAX=30
    endif
    do iter=1,IMAX

       call dfftw_r2c()
       fftcomplex2=fftcomplex

       ts=0D0

       do n=0,nexpand
          do j=1,nyi
             do i=1,nxi/2+1
                fftcomplex(i,j)=fftcomplex2(i,j)*expkz0kn(i,j,n)
             enddo
          enddo

          call dfftw_c2r()
          !fftreal is not normalized

          if(n .gt. 0) then
             do j=1,nyi
                do i=1,nxi
                   ts(i,j)=ts(i,j)+fftreal(i,j)*hnr_nfac(i,j,n)
                enddo
             enddo
          else
             do j=1,nyi
                do i=1,nxi
                   ts(i,j)=ts(i,j)+fftreal(i,j)*fftratio
                enddo
             enddo
          endif
       enddo

!       do j=1,nyi
!          do i=1,nxi
!             s(i,j)=2D0*u(i,j)-s(i,j)
!          enddo
!       enddo
!
!		fftreal=s

       fftreal=s+u-ts
	   s=fftreal

    enddo

    write(*,*)'                                ',maxval(abs(s))

  end subroutine equivsource


  subroutine upward()

    use mod_para
    use mod_grid
    use mod_bivar
    use mod_workspace
    use mod_fftw
    use mod_interp2

    call chz2bz()

    call griddata( x(nx1-1,ny1-1,nz1), y(nx1-1,ny1-1,nz1), chz(nx1-1,ny1-1,nz1-1), xi,yi, bzi)
	!write(11)nxi,nyi,deltah,xi,yi,zi,bzi
	!do while(.true.)
    call equivsource( bzi, bzis )
	!enddo

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!1
    !    call comp_bxi0k( ) 
    fftreal=bzis
    call dfftw_r2c()
    do j=1,nyi
       do i=1,nxi/2+1
          fftcomplex(i,j)=cmplx(0,1.D0)*kx(i)/kk(i,j)*fftcomplex(i,j)
       enddo
    enddo

    !    call comp_bxihk( ) 
    !    call comp_bxih( ) 

    fftcomplex2=fftcomplex
    bzi=0
    do n=0,nexpand
       do j=1,nyi
          do i=1,nxi/2+1
             fftcomplex(i,j)=fftcomplex2(i,j)*expkz0kn(i,j,n)
          enddo
       enddo
       call dfftw_c2r()
       if(n .gt. 0) then
          bzi=bzi+fftreal*hnr_nfacdh(:,:,n)
       else 
          bzi=bzi+fftreal*fftratio
       endif
    enddo


    !    call comp_bxh( ) 
    !call griddata2( x(nx1-1,ny1-1,nz1), y(nx1-1,ny1-1,nz1), chx(nx1-1,ny1-1,nz1-1), bzi )
    call interp2(bzi,chx(nx1-1,ny1-1,nz1-1))
	arrayview=chx(:,:,nz1-1)

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!1
    !    call comp_byi0k( ) 

    fftreal=bzis
    call dfftw_r2c()
    do j=1,nyi
       do i=1,nxi/2+1
          fftcomplex(i,j)=cmplx(0,1.D0)*ky(j)/kk(i,j)*fftcomplex(i,j)
       enddo
    enddo
    !    call comp_byihk( ) 
    !    call comp_byih( )  

    fftcomplex2=fftcomplex
    bzi=0
    do n=0,nexpand
       do j=1,nyi
          do i=1,nxi/2+1
             fftcomplex(i,j)=fftcomplex2(i,j)*expkz0kn(i,j,n)
          enddo
       enddo
       call dfftw_c2r()
       if(n .gt. 0) then
          bzi=bzi+fftreal*hnr_nfacdh(:,:,n)
       else 
          bzi=bzi+fftreal*fftratio
       endif

    enddo
    !    call comp_byh( ) 
    !call griddata2( x(nx1-1,ny1-1,nz1), y(nx1-1,ny1-1,nz1), chy(nx1-1,ny1-1,nz1-1), bzi )
    call interp2(bzi,chy(nx1-1,ny1-1,nz1-1))
	arrayview=chy(:,:,nz1-1)
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!1
    !    call comp_bzi0k( ) 
    fftreal=bzis
    call dfftw_r2c()

    !    call comp_bzihk( ) 
    !    call comp_bzih( ) 
    fftcomplex2=fftcomplex
    bzi=0
    do n=0,nexpand
       do j=1,nyi
          do i=1,nxi/2+1
             fftcomplex(i,j)=fftcomplex2(i,j)*expkz0kn(i,j,n)
          enddo
       enddo
       call dfftw_c2r()
       if(n .gt. 0) then
          bzi=bzi+fftreal*hnr_nfacdh(:,:,n)
       else 
          bzi=bzi+fftreal*fftratio
       endif
    enddo
    !    call comp_bzh( ) 
    !call griddata2( x(nx1-1,ny1-1,nz1), y(nx1-1,ny1-1,nz1), chz(nx1-1,ny1-1,nz1-1), bzi )
    call interp2(bzi,chz(nx1-1,ny1-1,nz1-1))
	arrayview=chz(:,:,nz1-1)
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!1
    call comp_chxy() 

	arrayview=chx(:,:,nz1-1)
	arrayview=chy(:,:,nz1-1)
	arrayview=chz(:,:,nz1-1)
	arrayview=chz(:,:,nz1)

  end subroutine upward


  subroutine comp_chxy()
    use mod_para
    use mod_grid
    use mod_bivar
    use mod_workspace
    use mod_fftw

    implicit none
    integer::i,j,k,tempk
    real(8)::t1,t2
    real(8)::a1x,a1y,a1z,a2x,a2y,a2z,a3x,a3y,a3z
    real(8)::ca1x,ca1y,ca1z,ca2x,ca2y,ca2z,ca3x,ca3y,ca3z

    k=nz1
    tempk=nz1-1
    do j=ny1,ny2
       do i=nx1,nx2
          a1x=x(i+1,j,k)-x(i,j,k)
          a1y=y(i+1,j,k)-y(i,j,k)
          a1z=z(i+1,j,k)-z(i,j,k)
          a2x=x(i,j+1,k)-x(i,j,k)
          a2y=y(i,j+1,k)-y(i,j,k)
          a2z=z(i,j+1,k)-z(i,j,k)
          a3x=x(i,j,k+1)-x(i,j,k)
          a3y=y(i,j,k+1)-y(i,j,k)
          a3z=z(i,j,k+1)-z(i,j,k)
          ca1x=a2y*a3z-a2z*a3y 
          ca1y=a2z*a3x-a2x*a3z
          ca1z=a2x*a3y-a2y*a3x
          ca2x=a3y*a1z-a3z*a1y 
          ca2y=a3z*a1x-a3x*a1z
          ca2z=a3x*a1y-a3y*a1x
          ca3x=a1y*a2z-a1z*a2y 
          ca3y=a1z*a2x-a1x*a2z
          ca3z=a1x*a2y-a1y*a2x

          t1=( ca1x*chx(i,j,tempk)+ca1y*chy(i,j,tempk)+ca1z*chz(i,j,tempk)  )/jac(i,j,k)
          t2=( ca2x*chx(i,j,tempk)+ca2y*chy(i,j,tempk)+ca2z*chz(i,j,tempk)  )/jac(i,j,k)
          chx(i,j,tempk)=t1
          chy(i,j,tempk)=t2
       enddo
    enddo

  end subroutine comp_chxy

  subroutine chz2bz()
    ! covert from curvilinear component to cartesian component, store in chz(:,:,nz-1)
    use mod_para
    use mod_grid
    use mod_bivar
    use mod_workspace
    use mod_fftw
    implicit none
    integer::tempk,i,j,k
    real(8)::a1x,a1y,a1z,a2x,a2y,a2z,a3x,a3y,a3z
    real(8)::ca1x,ca1y,ca1z,ca2x,ca2y,ca2z,ca3x,ca3y,ca3z
    k=nz1
    tempk=nz1-1
    do j=ny1,ny2
       do i=nx1,nx2
          a1x=x(i+1,j,k)-x(i,j,k)
          a1y=y(i+1,j,k)-y(i,j,k)
          a1z=z(i+1,j,k)-z(i,j,k)
          a2x=x(i,j+1,k)-x(i,j,k)
          a2y=y(i,j+1,k)-y(i,j,k)
          a2z=z(i,j+1,k)-z(i,j,k)
          a3x=x(i,j,k+1)-x(i,j,k)
          a3y=y(i,j,k+1)-y(i,j,k)
          a3z=z(i,j,k+1)-z(i,j,k)
          ca1x=a2y*a3z-a2z*a3y 
          ca1y=a2z*a3x-a2x*a3z
          ca1z=a2x*a3y-a2y*a3x
          ca2x=a3y*a1z-a3z*a1y 
          ca2y=a3z*a1x-a3x*a1z
          ca2z=a3x*a1y-a3y*a1x
          ca3x=a1y*a2z-a1z*a2y 
          ca3y=a1z*a2x-a1x*a2z
          ca3z=a1x*a2y-a1y*a2x
          chz(i,j,tempk)=(ca3x*sqrt(g11(i,j,k))*chx(i,j,k)+ca3y*sqrt(g22(i,j,k))*chy(i,j,k)+ca3z*sqrt(g33(i,j,k))*chz(i,j,k)) &
               /sqrt(ca3x*ca3x+ca3y*ca3y+ca3z*ca3z)

       enddo
    enddo

  endsubroutine chz2bz

end module mod_upward
